package com.lagou.mall.esflow.server.config;

import com.lagou.mall.esflow.server.ESFlowSystemException;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.function.Consumer;
import javax.annotation.Resource;
import javax.sql.DataSource;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.connection.ShardingConnection;
import org.springframework.stereotype.Component;

/**
 * @Author xiazhenyu
 * @Description
 * @Date 5:39 下午 2021/2/23
 **/
@Slf4j
@Component
public class ShardingJdbc {

    @Resource
    private DataSource shardingDataSource;


    public void getConnection(Consumer<Connection> connectionConsumer) {

        try (ShardingConnection connection = (ShardingConnection) shardingDataSource.getConnection()) {
            connectionConsumer.accept(connection);
        } catch (SQLException e) {
            throw ESFlowSystemException.error("create connection failed", e);
        }
    }


    public void newStatement(Consumer<Statement> statementConsumer) {
        newStatement(connection -> {
        }, statementConsumer);
    }

    public void newStatement(Consumer<Connection> connectionConsumer, Consumer<Statement> statementConsumer) {

        getConnection(connection -> {
            connectionConsumer.accept(connection);
            try (Statement statement = connection.createStatement()) {
                statementConsumer.accept(statement);
            } catch (SQLException e) {
                throw ESFlowSystemException.error("create statement failed", e);
            }
        });
    }


}
